package com.llmops.demo;

import com.llmops.core.Context;
import com.llmops.core.Edge;
import com.llmops.core.WorkflowGraph;
import com.llmops.core.node.*;

import java.util.Collections;
import java.util.Map;

public class Main1 {
    public static void main(String[] args) {
        WorkflowGraph workflow = new WorkflowGraph("1");
//        Context context = workflow.getContext();
//        context.setVariable("sessionId",6);

        workflow.addNode(new StartNode("dataInput","dataInput",Collections.singletonMap("input","小学生可以贷款吗")));
//        workflow.addNode(new ClassificationNode("llmClassificationNode",Collections.singletonMap("classifications", Arrays.asList("合规","不合规"))));
//
        workflow.addNode(new RagNode("rag","rag",Collections.singletonMap("input","{dataInput_output}")));

        String prompt = "您是一个全球信赖的专家问答系统。\n" +
                "始终使用提供的上下文信息回答问题，而不是先验知识。\n" +
                "必须遵守的规则：\n" +
                "1. 永远不要直接引用给定的上下文来回答，回复内容应尽可能简洁。\n" +
                "2. 避免使用“基于上下文，...”或“上下文信息...”或类似的表述。\n" +
                "3. 如果没有提供相关上下文，只需回答“文档信息不足，根据提供信息无法做出回答”。\n" +
                "4. 如果<data></data>标签内的内容包含网络链接（如http或https地址），则必须输出。当检测到链接为图片链接时，按照以下Markdown格式输出：%E9%93%BE%E6%8E%A5%E5%9C%B0%E5%9D%80；当检测到链接为非图片链接时，按照以下Markdown格式输出：描述；无需关心图片是否与用户的问题相关。例如：http://pa-mx.maipuat.msxf.test/resource?filePath=s3://rag-server/images/672daba5P5bQ.jpeg\n" +
                "5. 如果有多个图片链接，必须输出所有图片链接。\n" +
                "6.给出答案后，如果有必要请在最后根据用户的问题生成反问信息 如：\n" +
                "#用户：你好，请问你们提供个人消费贷款吗？\n" +
                "助手：您好！....。请问您对哪一种比较感兴趣呢？";

        workflow.addNode(new LLMNode("llm","llm", Map.of("systemPrompt",prompt,"userPrompt","背景知识如下[{rag_output}]，问题如下:{dataInput_output}")));

//        workflow.addNode(new OutputNode("out不合规", Collections.singletonMap("content", "请文明用语")));

        workflow.addNode(new EndNode("out合规","out合规", Collections.singletonMap("content", "{llm_output}")));

//        workflow.addNode(new OutputNode("result_exporterB", Collections.singletonMap("v", "B")));

        workflow.addEdge(new Edge("dataInput", "rag"));
//        workflow.addEdge(new Edge("llmClassificationNode", "rag", Edge.ConditionType.EQUAL,"合规"));
//        workflow.addEdge(new Edge("llmClassificationNode", "out不合规", Edge.ConditionType.EQUAL,"不合规"));
        workflow.addEdge(new Edge("rag", "llm"));
        workflow.addEdge(new Edge("llm", "out合规"));
//        workflow.addEdge(new Edge("llm_inferenceB", "result_exporterB"));

        workflow.validate();
        Context context = new Context();
        context.setVariable("sessionId",6);
//        workflow.execute();

        workflow.execute(context);

        System.out.println("Workflow execution completed");
    }
}